# coding=utf-8
# -*- coding: UTF-8 -*-
#
# This file is part of the kernelstudio package.
#
# (c) 2014-2025 zlin <admin@kernelstudio.com>
#
# For the full copyright and license information, please view the LICENSE file
# that was distributed with this source code.
import logging
import typing as t

import fitz
import numpy as np
from PIL import Image

from ksai.document.context import DocumentSection
from ksai.document.extractor.extractor import DocumentExtractorException, AbstractDocumentExtractor
from ksai.document.ocr.factory import ocr_factory
from ksai.persistence.storage.attachment import Attachment

logger = logging.getLogger(__name__)


def pdf_to_images(pdf_bytes):
    """将 PDF 字节流转换为图像列表"""
    images = []
    try:
        # 使用 PyMuPDF 打开 PDF 字节流
        doc = fitz.open(stream=pdf_bytes, filetype="pdf")
        for page_num in range(len(doc)):
            page = doc.load_page(page_num)
            # 设置缩放因子为 2 (提高分辨率)
            mat = fitz.Matrix(2.0, 2.0)
            pix = page.get_pixmap(matrix=mat, alpha=False)
            # 转换为 PIL Image
            img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
            images.append(img)
        doc.close()
    except Exception as e:
        logger.error(f"PDF 转换失败: {str(e)}")
        raise
    return images


class PdfDocumentExtractor(AbstractDocumentExtractor):

    def __init__(self):
        super().__init__()

    def do_extract(self, attachment: Attachment) -> t.List[DocumentSection]:
        try:
            # 转换为图像列表
            images = pdf_to_images(attachment.content)
            sections = []
            # 对每张图像进行 OCR
            index = 0
            for img in images:
                # 将 PIL Image 转换为 numpy array
                img_np = np.array(img)
                # 执行 OCR
                result = ocr_factory.executor.execute(img_np)
                sections.append(DocumentSection(
                    text=result,
                    page=index,
                ))
                index += 1
            return sections
        except Exception as e:
            raise DocumentExtractorException(f"Extract pdf file error: {e}") from e

    def extensions(self) -> t.List[str]:
        return ["pdf"]

    def mimes(self) -> t.List[str]:
        return ["application/pdf"]
